from pathlib import Path
from typing import Literal, TypeAlias

import addict
import gymnasium as gym
import numpy as np
import polars as pl
import toml
from numpy.typing import NDArray

EnvType: TypeAlias = Literal["train", "test", "valid"]
FloatArray: TypeAlias = NDArray[np.float32]

__all__ = ["HomeEnv"]


class HomeEnv(gym.Env):
    def __init__(
        self,
        config_path: Path | None = None,
        heter: FloatArray | None = None,
        seed: int | None = None,
        _type: EnvType = "train",
        auto_normalize: bool = False,
        v2g: bool = False,
    ) -> None:
        # 读取配置
        assert _type in EnvType.__args__
        self._type = _type
        self.auto_normalize = auto_normalize
        if config_path is None:
            config_path = Path(__file__).parent / "default.toml"
        self.config = addict.Dict(toml.load(config_path))

        # 设置持续时间
        if self._type == "test":
            self.day_duration = self.config.day_test
        elif self._type == "valid":
            self.day_duration = self.config.day_valid
        elif self._type == "train":
            assert 0 < self.config.day_train <= self.config.day_valid
            self.day_duration = self.config.day_train

        # 设置异质性
        self.setup_heter(heter)
        # 加载数据表
        self.read_table()
        # 获取数据范围供归一化使用
        self.data_scope = self.get_data_scope()
        # 记录环境信息
        self.meta_info = self.get_meta_info()

        # 设置动作空间和观测空间
        self.action_space = gym.spaces.Box(
            low=np.array([-self.p_ess_max, 0, -self.p_ev_max if v2g else 0]),
            high=np.array([self.p_ess_max, self.p_hvac_max, self.p_ev_max]),
        )
        if seed is not None:
            self.action_space.seed(seed)
        obs_low_bound = -np.array(
            [
                np.finfo(np.float32).max,
                np.finfo(np.float32).max,
                np.finfo(np.float32).max,
                np.finfo(np.float32).max,
                np.finfo(np.float32).max,
                np.finfo(np.float32).max,
                np.finfo(np.float32).max,
                np.finfo(np.float32).max,
                np.finfo(np.float32).max,
                np.finfo(np.float32).max,
            ]
        )
        self.observation_space = gym.spaces.Box(
            low=obs_low_bound, high=-obs_low_bound, dtype=np.float32
        )

    def read_table(self) -> None:
        deadline = self.config.day_test + self.config.day_valid + 1
        columns = ["{}:00".format(i) for i in range(24)]
        schema = dict(zip(columns, [pl.Float32] * len(columns)))
        self.solar_data_table = pl.read_csv(
            Path(__file__).parent / self.config.solar_data_path,
            columns=columns,
            dtypes=schema,
            n_rows=deadline,
        )
        self.load_data_table = pl.read_csv(
            Path(__file__).parent / self.config.load_data_path,
            columns=columns,
            dtypes=schema,
            n_rows=deadline,
        )
        self.temp_outdoor_data_table = pl.read_csv(
            Path(__file__).parent / self.config.temp_outdoor_data_path,
            columns=columns,
            dtypes=schema,
            n_rows=deadline,
        )
        self.price_data_table = pl.read_csv(
            Path(__file__).parent / self.config.price_data_path,
            columns=columns,
            dtypes=schema,
            n_rows=deadline,
        )
        self.ev_available_table = pl.read_csv(
            Path(__file__).parent / self.config.ev_available_path,
            columns=columns,
            dtypes=schema,
            n_rows=deadline,
        )
        self.ev_reduction_table = pl.read_csv(
            Path(__file__).parent / self.config.ev_reduction_path, n_rows=deadline
        )

    def setup_heter(self, heter: FloatArray | None) -> None:
        """异质参数设置"""
        if heter is not None:
            assert len(heter) <= 14
            while len(heter) < 14:
                heter = np.append(heter, 0.5)
        else:
            heter = np.full((14,), 0.50)
        self.eta_ess = (self.config.eta_ess - 0.03) + 0.06 * heter[0]
        self.ess_level_init = (self.config.ess_level_init - 0.42) + 0.84 * heter[1]
        self.ess_level_max = (self.config.ess_level_max - 1.8) + 3.6 * heter[2]
        self.ess_level_min = (self.config.ess_level_min - 0.18) + 0.36 * heter[3]
        self.ess_aging_cost = (self.config.ess_aging_cost - 0.0005) + 0.001 * heter[4]
        self.p_ess_max = (self.config.p_ess_max - 0.50) + 1.0 * heter[5]
        self.p_hvac_max = (self.config.p_hvac_max - 0.50) + 1.0 * heter[6]
        self.T_min = (self.config.T_min - 1.8) + 3.6 * heter[7]
        self.T_max = (self.config.T_max - 1.8) + 3.6 * heter[8]
        self.epsilon_hvac = (self.config.epsilon_hvac - 0.15) + 0.30 * heter[9]
        self.eta_hvac = (self.config.eta_hvac - 0.5) + 1.0 * heter[10]
        self.A_hvac = (self.config.A_hvac - 0.03) + 0.06 * heter[11]
        self.temp_indoor_init = (self.config.temp_indoor_init - 1.8) + 3.6 * heter[12]
        self.price_ratio = (self.config.price_ratio - 0.05) + 0.10 * heter[13]
        # 非异质参数
        self.beta_hvac = self.config.beta_hvac
        self.beta_ev = self.config.beta_ev
        self.ev_level_max = self.config.ev_level_max
        self.ev_level_min = self.config.ev_level_min
        self.eta_ev = self.config.eta_ev
        self.ev_level_init = self.config.ev_level_init
        self.ev_aging_cost = self.config.ev_aging_cost
        self.p_ev_max = self.config.p_ev_max

    def get_data_scope(self) -> FloatArray:
        """获取数据最大最小值, 供归一化使用"""
        return np.array(
            [
                [
                    self.solar_data_table.min().min_horizontal().item(),
                    self.load_data_table.min().min_horizontal().item(),
                    self.ess_level_min,
                    self.ev_level_min,
                    0,
                    0,
                    self.temp_outdoor_data_table.min().min_horizontal().item(),
                    self.T_min,
                    self.price_data_table.min().min_horizontal().item(),
                    0,
                ],
                [
                    self.solar_data_table.max().max_horizontal().item(),
                    self.load_data_table.max().max_horizontal().item(),
                    self.ess_level_max,
                    self.ev_level_max,
                    1,
                    23,
                    self.temp_outdoor_data_table.max().max_horizontal().item(),
                    self.T_max,
                    self.price_data_table.max().max_horizontal().item(),
                    23,
                ],
            ],
            dtype=np.float32,
        )

    def normalize_state(self, s: FloatArray) -> FloatArray:
        """归一化状态"""
        assert len(s.shape) == 1
        return (s - self.data_scope[0]) / (self.data_scope[1] - self.data_scope[0])

    def denormalize_state(self, s: FloatArray) -> FloatArray:
        """恢复被归一化的状态"""
        assert len(s.shape) == 1
        return s * (self.data_scope[1] - self.data_scope[0]) + self.data_scope[0]

    def clip_action(self, a: FloatArray) -> FloatArray:
        """限制动作"""
        p_ess, p_hvac, p_ev = a
        (
            p_solar,
            p_load,
            ess_level,
            ev_level,
            ev_ok,
            ev_departure,
            temp_outdoor,
            temp_indoor,
            price,
            _,
        ) = self.state

        # 限制ess
        if p_ess >= 0:
            p_ess = np.clip(
                p_ess,
                0,
                min(
                    (self.ess_level_max - ess_level) / self.eta_ess,
                    self.p_ess_max,
                ),
            )
        else:
            p_ess = -np.clip(
                -p_ess,
                0,
                min(
                    (ess_level - self.ess_level_min) * self.eta_ess,
                    self.p_ess_max,
                ),
            )

        # 限制hvac
        if temp_indoor <= self.T_min:
            p_hvac = 0
        if temp_indoor > self.T_max:
            p_hvac = np.clip(p_hvac, 0.1, self.p_hvac_max)

        # 限制ev
        if ev_ok == 0:
            p_ev = 0
        elif p_ev > 0:
            p_ev = np.clip(
                p_ev,
                0,
                min(
                    (self.ev_level_max - ev_level) / self.eta_ev,
                    self.p_ev_max,
                ),
            )
        else:
            p_ev = -np.clip(
                -p_ev,
                0,
                min(
                    (ev_level - self.ev_level_min) * self.eta_ev,
                    self.p_ev_max,
                ),
            )
        return np.array([p_ess, p_hvac, p_ev], np.float32)

    def get_ev_departure(self) -> int:
        """获取ev离开时间"""
        day, hour = self.day, self.hour
        cursor = 0
        if self.ev_available_table.item(day, hour) == 0:
            return -1
        while self.ev_available_table.item(day, hour) == 1:
            cursor += 1
            hour += 1
            if hour > 23:
                hour = 0
                day += 1
        return (day - self.day) * 24 + hour - self.hour

    def step(self, a: FloatArray) -> tuple[FloatArray, float, bool, bool, dict]:
        """环境迭代"""
        (
            p_solar,
            p_load,
            ess_level,
            ev_level,
            ev_ok,
            ev_departure,
            temp_outdoor,
            temp_indoor,
            price,
            _,
        ) = self.state
        p_ess, p_hvac, p_ev = self.clip_action(a)

        # 计算ESS下一时隙的储能
        ess_level_next = ess_level + (
            p_ess * self.eta_ess if p_ess > 0 else p_ess / self.eta_ess
        )
        ess_level_next = np.clip(ess_level_next, self.ess_level_min, self.ess_level_max)

        # 计算ev下一时隙的储能
        if ev_ok != 0:
            ev_level_next = ev_level + (
                p_ev * self.eta_ev if p_ev > 0 else p_ev / self.eta_ev
            )
            ev_level_next = np.clip(ev_level_next, self.ev_level_min, self.ev_level_max)
        else:
            ev_level_next = -1

        # 计算下一时隙的室内温度
        temp_indoor_next = self.epsilon_hvac * temp_indoor + (1 - self.epsilon_hvac) * (
            temp_outdoor - self.eta_hvac * p_hvac / self.A_hvac
        )

        # 计算向主电网的购电/售电量
        p_grid = p_load + p_ess + p_hvac + p_ev - p_solar

        # 计算reward
        c1 = p_grid * price
        if p_grid < 0:
            c1 *= self.price_ratio
        c2 = abs(p_ess) * self.ess_aging_cost
        c3 = float(
            max(0, self.T_min - temp_indoor_next)
            + max(0, temp_indoor_next - self.T_max)
        )
        r = -self.beta_hvac * (c1 + c2) - c3

        # 更新时间设置
        self.n_step += 1
        self.hour += 1
        if self.hour > 23:
            self.hour = 0
            self.day += 1

        # 检测到0-1上升沿时重置电量, 根据离家时电量进行衰减
        ev_departure_next = self.get_ev_departure()
        ev_ok_next = ev_departure_next > 0
        if ev_ok == 0 and ev_ok_next != 0:
            ev_level_next = max(
                self.ev_level_bkp - self.ev_reduction_table.item(self.day, 0),
                self.ev_level_min,
            )

        # 到达depature时刻，如果电池未充满，增加一个惩罚项
        c4 = 0
        if ev_ok != 0 and ev_ok_next == 0:
            c4 = self.ev_level_max - ev_level_next
            r -= self.beta_ev * c4
            self.ev_level_bkp = ev_level_next

        # 更新内状态，内部状态始终不会被归一化
        self.state = np.array(
            [
                self.solar_data_table.item(self.day, self.hour),
                self.load_data_table.item(self.day, self.hour),
                ess_level_next,
                ev_level_next,
                ev_ok_next,
                ev_departure_next,
                self.temp_outdoor_data_table.item(self.day, self.hour),
                temp_indoor_next,
                self.price_data_table.item(self.day, self.hour),
                self.hour,
            ],
            dtype=np.float32,
        )

        # 归一化输出状态
        if self.auto_normalize is True:
            state = self.normalize_state(self.state)
        else:
            state = self.state.copy()

        # 判断是否结束
        t1 = False
        t2 = self.n_step >= self.day_duration * 24
        info = {
            "c1": c1,
            "c2": c2,
            "c3": c3,
            "c4": c4,
            "p_grid": p_grid,
            "a": [p_ess, p_hvac, p_ev],
        }
        return state, r, t1, t2, info

    def reset(
        self, *, seed: int | None = None, options: dict | None = None
    ) -> tuple[FloatArray, dict]:
        """重置环境"""
        if self._type == "test":
            self.day = 61
        elif self._type == "valid":
            self.day = 0
        elif self._type == "train":
            self.day = np.random.randint(
                0, self.config.day_valid + 1 - self.day_duration
            )
        self.n_step = 0
        self.hour = 0
        self.state = np.array(
            [
                self.solar_data_table.item(self.day, 0),
                self.load_data_table.item(self.day, 0),
                self.ess_level_init,
                self.ev_level_init,
                self.ev_available_table.item(self.day, 0),
                self.get_ev_departure(),
                self.temp_outdoor_data_table.item(self.day, 0),
                self.temp_indoor_init,
                self.price_data_table.item(self.day, 0),
                0,
            ],
            dtype=np.float32,
        )
        info = {"day": self.day, "hour": self.hour}
        if self.auto_normalize is True:
            state = self.normalize_state(self.state)
        else:
            state = self.state.copy()
        return state, info

    def get_meta_info(self) -> dict[str, float]:
        """获取环境元信息, 主要是异质参数"""
        info = {}
        heters = [
            "eta_ess",
            "ess_level_init",
            "ess_level_max",
            "ess_level_min",
            "ess_aging_cost",
            "p_ess_max",
            "p_hvac_max",
            "T_min",
            "T_max",
            "epsilon_hvac",
            "eta_hvac",
            "A_hvac",
            "temp_indoor_init",
            "price_ratio",
        ]
        for k in heters:
            info.update({k: getattr(self, k)})
        return info
